import websockets
import asyncio
import numpy as np
from funasr import AutoModel
import json
import logging
from typing import List
import io
import base64
import soundfile
import subprocess

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class VADServer:
    def __init__(self):
        print("初始化")
        self.model = AutoModel(model="/root/yuehu/assets/speech_fsmn_vad_zh-cn-16k-common-pytorch", device='cuda:1', disable_update=True)
        # self.model = AutoModel(model="fsmn-vad")
        self.chunk_size = 60  # ms
        self.sample_rate = 16000  # 采样率
        self.cache = {}
        self.vad_res = []
        self.detectedId = 0

    def check_vad(self, vad_list: List[List[float]]) -> float:
        """检查VAD结果"""
        if len(vad_list) > 0 and len(vad_list) % 2 == 0:
            last_two = vad_list[-2:]
            max_values = [max(sublist) for sublist in last_two]
            return abs(max_values[1] - max_values[0])
        return 0

    async def process_audio(self, audio_bytes: bytes, request_id: int) -> dict:
        """处理音频数据并返回VAD检测结果"""
        try:
            audio_stream = io.BytesIO(audio_bytes)
            speech, sample_rate = soundfile.read(audio_stream)
            chunk_stride = int(self.chunk_size * sample_rate / 1000)
            total_chunk_num = int((len(speech) - 1) / chunk_stride + 1)
            print(f"speech len {len(speech)},sample_rate {sample_rate},chunk_stride {chunk_stride},total_chunk_num {total_chunk_num}")

            # 逐块处理音频
            for i in range(total_chunk_num):
                speech_chunk = speech[i * chunk_stride:(i + 1) * chunk_stride]
                # is_final = i == total_chunk_num - 1
                is_final = False
                res = self.model.generate(
                    input=speech_chunk,
                    # output_dir='output_dir/',
                    cache=self.cache,
                    is_final=is_final,
                    chunk_size=self.chunk_size,
                    max_end_silence_time=500,
                    max_start_silence_time=1000
                )

                if len(res) == 0:
                    continue
                value = res[0]['value']
                if len(value) == 0:
                    continue
                self.vad_res.append(res[0]['value'][0])

                # # 表示至少识别 1 秒的长度才认为需要结束.
                # if self.check_vad(self.vad_res) > 1000:
                #     print(f"vad_res: {self.vad_res}")
                #     break
                if self.check_vad(self.vad_res) > 1000:
                    self.detectedId = request_id

            return {
                'status': 'success',
                'vad_results': self.vad_res,
                'isDetected': self.check_vad(self.vad_res) >= 1000,
                'detectedId': self.detectedId,
                'requestId': request_id
            }

        except Exception as e:
            logger.error(f"Error processing audio: {e}")
            return {
                'status': 'error',
                'error': str(e)
            }

def detect_audio_format(data):
    """
    检测音频格式

    参数:
        data (bytes): 音频数据

    返回:
        str: 'wav' 或 'opus' 或 'unknown'
    """
    # 检查 WAV 头部标识
    # 检查 Ogg 容器中的 Opus
    if len(data) > 4 and data[:4] == b'OggS':
        # 通常 Opus 在 Ogg 容器中的标识在后面
        # 简单搜索特征字节
        if b'OpusHead' in data[:100]:  # 检查前100字节
            #print("opus 2")
            return 'opus'

    # Opus 文件通常以 'OpusHead' 标记开始
    # 检查标准的 Opus 头部
    if len(data) > 8 and data[:8] == b'OpusHead':
        #print("opus 1")
        return 'opus'

    if len(data) > 12 and data[:4] == b'RIFF' and data[8:12] == b'WAVE':
        #print("wav")
        return 'wav'

    return 'unknown'

def convert_opus_to_wav(opus_data: bytes) -> bytes:
    format_type = detect_audio_format(opus_data)
    if format_type == 'opus':
        try:
            # Prepare FFmpeg command
            ffmpeg_cmd = [
                'ffmpeg',
                '-i', 'pipe:0',  # Input from pipe
                '-ac', '1',  # Set to mono
                '-ar', '16000',  # Set sample rate to 16kHz
                '-f', 'wav',  # Output format
                'pipe:1'  # Output to pipe
            ]

            # Start FFmpeg process
            process = subprocess.Popen(
                ffmpeg_cmd,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )

            # Feed opus data and get WAV output
            stdout_data, stderr_data = process.communicate(input=opus_data)

            # Check if process was successful
            if process.returncode != 0:
                error_msg = f"FFmpeg process failed with exit code {process.returncode}"
                print(error_msg)
                return opus_data

            return stdout_data

        except Exception as e:
            print(f"转换失败: {str(e)}")
            return opus_data
    else:
        return opus_data

async def handle_client(websocket):
    """处理WebSocket客户端连接"""
    # 初始化VAD服务器
    vad_server = VADServer()
    print("进入")
    try:
        async for message in websocket:
            try:
                # 解析接收到的JSON消息
                data = json.loads(message)
                # 将十六进制字符串转换回字节数组
                audio_bytes = base64.b64decode(data['audio_data'], validate=True)
                request_id = data['request_id']
                decoded_audio = convert_opus_to_wav(audio_bytes)
                # 处理音频数据
                result = await vad_server.process_audio(decoded_audio, request_id)

                # 发送结果回客户端
                await websocket.send(json.dumps({
                    'request_id': request_id,
                    'result': result
                }))

            except json.JSONDecodeError:
                await websocket.send(json.dumps({
                    'status': 'error',
                    'error': 'Invalid JSON format'
                }))
            except Exception as e:
                await websocket.send(json.dumps({
                    'status': 'error',
                    'error': str(e)
                }))
    except websockets.exceptions.ConnectionClosed:
        logger.info("Client connection closed")


async def main():
    # 启动WebSocket服务器
    server = await websockets.serve(
        lambda ws: handle_client(ws),  # 移除了path参数
        "0.0.0.0",
        # "localhost",
        8765,
        max_size=10485760  # 10MB最大消息大小
    )

    logger.info("VAD WebSocket Server started on ws://localhost:8765")
    await server.wait_closed()


if __name__ == "__main__":
    asyncio.run(main())
